## Two-variable interaction plots in R
## Anton Strezhnev 
## 06/17/2013
## with modifications by Jonathan Homola

interaction_plot_continuous <- function(model, effect, moderator, interaction, varcov="default", minimum="min", maximum="max", ymin="min", ymax="max", incr="default", num_points = 50, conf=.95, mean=FALSE, median=FALSE, alph=80, rugplot=T, histogram=T, pid=F, ideol=F, title="Marginal effects plot", xlabel="Value of moderator", ylabel="Estimated marginal coefficient", ...){
  
  # Define a function to make colors transparent
  makeTransparent<-function(someColor, alpha=alph){
    newColor<-col2rgb(someColor)
    apply(newColor, 2, function(curcoldata){rgb(red=curcoldata[1], green=curcoldata[2],
                                                blue=curcoldata[3],alpha=alpha, maxColorValue=255)})
  }
  
  # Extract Variance Covariance matrix
  if (varcov == "default"){
    covMat = vcov(model)
  }else{
    covMat = varcov
  }
  
  # Extract the data frame of the model
  mod_frame = model.frame(model)
  
  # Get coefficients of variables
  beta_1 = model$coefficients[[effect]]
  beta_3 = model$coefficients[[interaction]]
  
  # Set range of the moderator variable
  # Minimum
  if (minimum == "min"){
    min_val = min(mod_frame[[moderator]])
  }else{
    min_val = minimum
  }
  # Maximum
  if (maximum == "max"){
    max_val = max(mod_frame[[moderator]])
  }else{
    max_val = maximum
  }
  
  # Check if minimum smaller than maximum
  if (min_val > max_val){
    stop("Error: Minimum moderator value greater than maximum value.")
  }
  
  # Determine intervals between values of the moderator
  if (incr == "default"){
    increment = (max_val - min_val)/(num_points - 1)
  }else{
    increment = incr
  }
  
  # Create list of moderator values at which marginal effect is evaluated
  x_2 <- seq(from=min_val, to=max_val, by=increment)
  
  # Compute marginal effects
  delta_1 = beta_1 + beta_3*x_2
  
  # Compute variances
  var_1 = covMat[effect,effect] + (x_2^2)*covMat[interaction, interaction] + 2*x_2*covMat[effect, interaction]
  
  # Standard errors
  se_1 = sqrt(var_1)
  
  # Upper and lower confidence bounds
  z_score = qnorm(1 - ((1 - conf)/2))
  upper_bound = delta_1 + z_score*se_1
  lower_bound = delta_1 - z_score*se_1
  
  # Determine the bounds of the graphing area
  if (ymax == "max"){
    max_y = max(upper_bound)
  }else{
    max_y = ymax
  }
  if (ymin == "min"){
    min_y = min(lower_bound)
  }else{
    min_y = ymin
  }
  
  # Make the histogram color
  hist_col = makeTransparent("grey")
  
  # Initialize plotting window
  if (pid) {
    plot(x=c(), y=c(), ylim=c(min_y, max_y), xlim=c(min_val, max_val), xlab=xlabel, ylab=ylabel, main=title, xaxt="n")
    axis(side=1, at=c(0.25, 3, 5.75), labels=c("Strong Democrat", "Independent", "Strong Republican"), las=1)
  }else if (ideol) {
    plot(x=c(), y=c(), ylim=c(min_y, max_y), xlim=c(min_val, max_val), xlab=xlabel, ylab=ylabel, main=title, xaxt="n")
    axis(side=1, at=c(0.25, 3, 5.75), labels=c("Very Liberal", "Moderate", "Very Conservative"), las=1)
  }else{  
  plot(x=c(), y=c(), ylim=c(min_y, max_y), xlim=c(min_val, max_val), xlab=xlabel, ylab=ylabel, main=title)
  }
  
  # Plot estimated effects
  lines(y=delta_1, x=x_2)
  lines(y=upper_bound, x=x_2, lty=2)
  lines(y=lower_bound, x=x_2, lty=2)
  
  # Add a dashed horizontal line for zero
  abline(h=0, lty=3)
  
  # Add a vertical line at the mean
  if (mean){
    abline(v = mean(mod_frame[[moderator]]), lty=2, col="red")
  }
  
  # Add a vertical line at the median
  if (median){
    abline(v = median(mod_frame[[moderator]]), lty=3, col="blue")
  }
  
  # Add Rug plot
  if (rugplot){
    rug(mod_frame[[moderator]])
  }
  # Add Histogram (Histogram only plots when minimum and maximum are the min/max of the moderator)
  if (histogram & minimum=="min" & maximum=="max"){
    par(new=T)
    hist(mod_frame[[moderator]], axes=F, xlab="", ylab="",main="", border=hist_col, col=hist_col)
  }
  # Add special 7 point PID Histogram
  if (pid & minimum=="min" & maximum=="max"){
    par(new=T)
    hist(mod_frame[[moderator]]+0.25, axes=F, xlab="", ylab="",main="", border=hist_col, col=hist_col)
  }
  # Add special 7 point Ideology Histogram
  if (ideol & minimum=="min" & maximum=="max"){
    par(new=T)
    hist(mod_frame[[moderator]]+0.25, axes=F, xlab="", ylab="",main="", border=hist_col, col=hist_col)
  }
}


interaction_plot_binary <- function(model, effect, moderator, interaction, varcov="default", conf=.95, title="Marginal effects plot", xlabel="Value of moderator", ylabel="Estimated marginal coefficient", ymax="max", ymin="min", factor_labels=c(0,1)){
  
  # Extract Variance Covariance matrix
  if (varcov == "default"){
    covMat = vcov(model)
  }else{
    covMat = varcov
  }
  
  # Extract the data frame of the model
  mod_frame = model.frame(model)
  
  # Get coefficients of variables
  beta_1 = model$coefficients[[effect]]
  beta_3 = model$coefficients[[interaction]]
  
  # Create list of moderator values at which marginal effect is evaluated
  x_2 <- c(0,1)
  
  # Compute marginal effects
  delta_1 = beta_1 + beta_3*x_2
  
  # Compute variances
  var_1 = covMat[effect,effect] + (x_2^2)*covMat[interaction, interaction] + 2*x_2*covMat[effect, interaction]
  
  # Standard errors
  se_1 = sqrt(var_1)
  
  # Upper and lower confidence bounds
  z_score = qnorm(1 - ((1 - conf)/2))
  upper_bound = delta_1 + z_score*se_1
  lower_bound = delta_1 - z_score*se_1
  
  # Determine the bounds of the graphing area
  if (ymax == "max"){
    max_y = max(upper_bound)
  }else{
    max_y = ymax
  }
  if (ymin == "min"){
    min_y = min(lower_bound)
  }else{
    min_y = ymin
  }
  
  # Initialize plotting window
  plot(x=c(), y=c(), ylim=c(-0.75, 1.75), xlim=c(min_y, max_y), xlab=xlabel, ylab=ylabel, main=title, yaxt="n")
  
  # Plot points of estimated effects
  points(x=delta_1, y=x_2, pch=16)
  
  # Plot lines of confidence intervals
  lines(x=c(upper_bound[1], lower_bound[1]), y=c(x_2[1], x_2[1]), lty=1)
  #points(x=c(upper_bound[1], lower_bound[1]), y=c(x_2[1], x_2[1]), pch=c(25,24), bg="black")
  lines(x=c(upper_bound[2], lower_bound[2]), y=c(x_2[2], x_2[2]), lty=1)
  #points(x=c(upper_bound[2], lower_bound[2]), y=c(x_2[2], x_2[2]), pch=c(25,24), bg="black")
  
  # Label the axis
  #axis(side=1, at=c(0,1), labels=factor_labels)
  axis(side=2, at=c(0,1), labels=factor_labels, las=1)
  
  # Add a dashed horizontal line for zero
  abline(v=0, lty=1, col="red")
  
}


## interaction plot binary3 for left voters, right voters, non-voters in yougov data
interaction_plot_binary3 <- function(model, effect, moderator1, interaction1, moderator2, interaction2, varcov="default", conf=.95, title="Marginal effects plot", xlabel="Value of moderator", ylabel="Estimated marginal coefficient", ymax="max", ymin="min", factor_labels=c(0,1)){
  
  # Extract Variance Covariance matrix
  if (varcov == "default"){
    covMat = vcov(model)
  }else{
    covMat = varcov
  }
  
  # Extract the data frame of the model
  mod_frame = model.frame(model)
  
  # Get coefficients of variables
  beta_1 = model$coefficients[[effect]]
  beta_3a = model$coefficients[[interaction1]]
  beta_3b = model$coefficients[[interaction2]]
  
  # Create list of moderator values at which marginal effect is evaluated
  x_2 <- c(0,1)
  
  # Compute marginal effects
  delta_1a = beta_1 + beta_3a*x_2
  delta_1b = beta_1 + beta_3b*x_2
  
  # Compute variances
  var_1a = covMat[effect,effect] + (x_2^2)*covMat[interaction1, interaction1] + 2*x_2*covMat[effect, interaction1]
  var_1b = covMat[effect,effect] + (x_2^2)*covMat[interaction2, interaction2] + 2*x_2*covMat[effect, interaction2]
  
  # Standard errors
  se_1a = sqrt(var_1a)
  se_1b = sqrt(var_1b)
  
  # Upper and lower confidence bounds
  z_score = qnorm(1 - ((1 - conf)/2))
  upper_bounda = delta_1a + z_score*se_1a
  lower_bounda = delta_1a - z_score*se_1a
  upper_boundb = delta_1b + z_score*se_1b
  lower_boundb = delta_1b - z_score*se_1b
  
  # Determine the bounds of the graphing area
  if (ymax == "max"){
    max_y = max(upper_bounda)
  }else{
    max_y = ymax
  }
  if (ymin == "min"){
    min_y = min(lower_bounda)
  }else{
    min_y = ymin
  }
  
  # Initialize plotting window
  plot(x=c(), y=c(), ylim=c(-0.75, 1.75), xlim=c(min_y, max_y), xlab=xlabel, ylab=ylabel, main=title, yaxt="n")
  
  # Plot points of estimated effects
  x_2b = c(0, 0.50)
  x_2a = c(0, 1)
  points(x=delta_1a, y=x_2a, pch=16)
  points(x=delta_1b, y=x_2b, pch=16)
  
  # Plot lines of confidence intervals
  lines(x=c(upper_bounda[1], lower_bounda[1]), y=c(x_2a[1], x_2a[1]), lty=1)
  lines(x=c(upper_bounda[2], lower_bounda[2]), y=c(x_2a[2], x_2a[2]), lty=1)
  lines(x=c(upper_boundb[1], lower_boundb[1]), y=c(x_2b[1], x_2b[1]), lty=1)
  lines(x=c(upper_boundb[2], lower_boundb[2]), y=c(x_2b[2], x_2b[2]), lty=1)
  
  # Label the axis
  axis(side=2, at=c(0, 0.5, 1), labels=c("Leftist", "Nonvoter", "Rightist"), las=1)
  
  # Add a dashed horizontal line for zero
  abline(v=0, lty=1, col="red")
  
}